import asyncio
from autogen_ext.models.openai import OpenAIChatCompletionClient
from autogen_core.models import ModelFamily
from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_agentchat.conditions import TextMentionTermination
from autogen_agentchat.ui import Console
from utils.get_llm_config import get_llm_config


async def main():
    llm_config = get_llm_config()

    # 配置模型客户端
    # model_client = OpenAIChatCompletionClient(
    #     **llm_config,
    #     model_info={
    #         "vision": True,
    #         "function_calling": True,
    #         "json_output": True,
    #         "family": ModelFamily.UNKNOWN,
    #         "structured_output": True,
    #     }
    # )
    model_client = OpenAIChatCompletionClient(
        **llm_config,
        model_info={
            "vision": False,
            "function_calling": False,
            "json_output": False,
            "family": ModelFamily.UNKNOWN,
            "structured_output": True,
        },
        timeout=300,
        max_retries=3,
    )

    from pydantic import BaseModel

    class PRDOutputModel(BaseModel):
        title: str
        description: str
        requirements: str
        technical_architecture: str
        user_interface: str

    def get_doc(doc_path: str):
        """读取指定路径, 转为str"""
        with open(doc_path, "r", encoding="utf-8") as f:
            return f.read()
    example_prd_doc = get_doc("docs/example_prd.txt")

    prd_system_message = """你是产品需求分析师。根据用户需求生成完整的PRD文档.
<example_prd_doc>
{example_prd_doc}
</example_prd_doc>
按照示例, 输出专业的markdown格式PRD文档。""".format(example_prd_doc=example_prd_doc)
    # 创建三个专业智能体
    prd_writer = AssistantAgent(
        name="prd_writer",
        model_client=model_client,
        system_message=prd_system_message,
    )

    # result = await prd_writer.run(task="介绍一下你自己")

    # project_goal = "设计一个自动生成带货文章的产品."
    project_goal = ""

    real_prd_doc = get_doc("docs/example_prd_result.txt")

    # prd_result = await prd_writer.run(task=project_goal)
    # real_prd_doc = prd_result.messages[-1].content

    # print(f"--- tt.now: {tt.now()} ----- start real_prd_doc\n")  # type: ignore
    # print(real_prd_doc)  # type: ignore
    # print(f"--- tt.now: {tt.now()} ===== end real_prd_doc\n")  # type: ignore

    from prompts.prompts import BASE_PLANNING_PROMPT
    from bdtime import tt

    planning_system_message = BASE_PLANNING_PROMPT
    planning_system_message = planning_system_message.replace("{project_goal}", project_goal)
    planning_system_message = planning_system_message.replace("{requirement_document}", real_prd_doc)

    # print(planning_system_message)
    print(f"--- tt.now: {tt.now()} ----- planning_system_message\n")
    task_manager = AssistantAgent(
        name="task_manager",
        model_client=model_client,
        system_message=planning_system_message
    )
    print(f"--- tt.now: {tt.now()} --- start planning_json_str\n")

    # --- history data
    _planning_result_eg_file_path = "docs/planning_1.txt"
    planning_json_str = get_doc(_planning_result_eg_file_path)

    # planning_result = await task_manager.run(task=project_goal)
    # planning_json_str = planning_result.messages[-1].content
    #
    # with open(_planning_result_eg_file_path, 'w+', encoding='utf-8') as f:
    #     f.write(planning_json_str)

    # print(planning_json_str)
    # print(f"--- tt.now: {tt.now()} ===== end planning_json_str\n")

    reviewer_system_message = """你是质量审核专家。审核PRD文档和任务拆解的完整性、合理性，提供改进建议。
<评估标准参考>
{example_evaluation_rule}
</评估标准参考>
若判定任务拆分结果质量合格, 不需要修改了, 则最后输出 'COMPLETE' 表示工作流结束。"""
    example_evaluation_rule = get_doc("docs/规划评估标准参考.txt")
    reviewer_system_message = reviewer_system_message.replace("{example_evaluation_rule}", example_evaluation_rule)
    reviewer = AssistantAgent(
        name="reviewer",
        model_client=model_client,
        system_message=reviewer_system_message
    )

    reviewer_result = await reviewer.run(task=planning_json_str)
    reviewer_json_str = reviewer_result.messages[-1].content

    print(f"--- tt.now: {tt.now()} --- start reviewer_json_str\n")
    print(reviewer_json_str)
    print(f"--- tt.now: {tt.now()} ===== end reviewer_json_str\n")

    # --- history data
    _reviewer_result_eg_file_path = "docs/reviewer_1.txt"
    # reviewer_json_str = get_doc(_reviewer_result_eg_file_path)

    with open(_reviewer_result_eg_file_path, 'w+', encoding='utf-8') as f:
        print(f"--- save to _reviewer_result_eg_file_path: {_reviewer_result_eg_file_path}")
        f.write(reviewer_json_str)
    exit()

    # 创建工作流团队
    termination = TextMentionTermination("COMPLETE")

    # members = [prd_writer, task_manager, reviewer]
    members = [task_manager, reviewer]
    team = RoundRobinGroupChat(members, termination_condition=termination)

    # 用户需求
    user_requirement = input("请输入您的开发需求: ")

    task_prompt = f"""
用户开发需求：{user_requirement}

请按顺序完成：
1. PRD_WRITER: 生成完整PRD文档
2. TASK_MANAGER: 拆解为开发子任务  
3. REVIEWER: 审核并确认完成
"""

    # 运行工作流
    print("🚀 开始生成PRD和拆解任务...")
    await Console(team.run_stream(task=task_prompt))

if __name__ == "__main__":
    asyncio.run(main())
